'''
Created on Jan 13, 2015

@author: dewey
'''

import sys
import os
import subprocess
import numpy as np #do not remove! needed for eval() on custom error functions
import functools
import traceback

from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
from PyQt5.QtGui import QIcon, QColor, QBrush

import pystrat.project as pystratproj
import pystrat.datamodel as pystratmodel
import pystrat.dataio as pystratIO
import pystrat.units as units
import pystrat.dataprocessing as dataprocessing


from ui_gen.Ui_pystrat_main import Ui_PyStratMainWindow
from ui.autochart_dialog import PyStratAutoChartDialog

DIRTY = False
DEFAULT_PREFS = {}

DEFAULT_PREFS["test_key"] = "test value"

PREF_ERROR = "error"
ERROR_STANDARD_DEVIATION = "stdev"
ERROR_STANDARD_ERROR = "stdev/np.sqrt(n)"
ERROR_RELATIVE_STDEV_ERROR = "stdev/value"

DEFAULT_PREFS[PREF_ERROR] = ERROR_STANDARD_DEVIATION

PREF_RECENT_FILES = "recent_files"
DEFAULT_PREFS[PREF_RECENT_FILES] = []

PREF_NUMBER_FORMAT_VALUE = "number_format_value"
PREF_NUMBER_FORMAT_ERROR = "number_format_error"

PREF_LAST_OPEN_FILE = "last_open_file"

DEFAULT_PREFS[PREF_NUMBER_FORMAT_VALUE] = "%0.3f"
DEFAULT_PREFS[PREF_NUMBER_FORMAT_ERROR] = "%0.3f"

PREF_REPLICATE_BG_COLOR = "replicate_bg_color"
DEFAULT_PREFS[PREF_REPLICATE_BG_COLOR] = "#D2E3F3"

PREFS = DEFAULT_PREFS

ROLE_GLOBAL = "role_global"

TYPES_SAVE = "PyStrat Project File (*.pdt)\nComma Separated Values File (*.csv)"
TYPES_IMPORT = "PyStrat Types (*.pdt; *.csv)\nPyStrat Project File (*.pdt)\nComma Separated Values (*.csv)\nAll Files (*.*)"
TYPES_ALL = "All Files (*.*)"


#custom project keys
KEY_TREE_EXPANDED = "pystrat_tree_expanded"

def calculate_error(value, n, stdev):
    return eval(PREFS[PREF_ERROR])

def format_value(value):
    if not value is None:
        return PREFS[PREF_NUMBER_FORMAT_VALUE] % value
    else:
        return None

def format_error(error_value):
    if not error_value is None:
        return PREFS[PREF_NUMBER_FORMAT_ERROR] % error_value
    else:
        return None

def shellquote(s):
    return "'" + s.replace("'", "'\\''") + "'"

def launchfile(filename):
    if sys.platform.startswith("linux"):
        subprocess.call(["xdg-open", filename])
    elif sys.platform == 'win32':
        os.system('start %s' % shellquote(filename))
    elif sys.platform == 'darwin':
        os.system('open %s' % shellquote(filename))

def log(message):
    print("log: %s" % message)
    
def loc_id(vdata):
    return vdata.metadata(pystratproj.METADATA_LOCATION_ID)

class PyStratMainWindow(QMainWindow, Ui_PyStratMainWindow):
    
    def __init__(self, parent=None):
        super(PyStratMainWindow, self).__init__(parent)
        self.LOADING = True
        self.setupUi(self)
        self.restore_window_settings()
        self.restore_preferences()
        self.refresh_recent_files()
        
        self.TEMPFILES = []
        self.PROJECT = pystratproj.PyStratProject()
        self.PROJECT_DIRTY = False
        self.PROJECT_FILE = None
        
        self.autochartDialog = PyStratAutoChartDialog(self)
        self.autochartDialog.set_error_function(calculate_error)
        self.autochartDialog.set_prefs(PREFS)
        self.autochartDialog.set_tempfiles_list(self.TEMPFILES)
        
        self.setup_trees()
        self.setup_data_table()
        
        self.LOADING = False
        
        if PREF_LAST_OPEN_FILE in PREFS:
            self.load_file(PREFS[PREF_LAST_OPEN_FILE])
        
    
    #===============OPEN, SAVE, IMPORT, SET PROJECT=================#
    
    @pyqtSlot()
    def on_actionOpen_triggered(self):
        if self.PROJECT_DIRTY:
            if self.ask_save_changes():
                fname = self.getfile(TYPES_IMPORT, ROLE_GLOBAL)
                self.load_file(fname)
        else:
            fname = self.getfile(TYPES_IMPORT, ROLE_GLOBAL)
            self.load_file(fname)
    
    @pyqtSlot()
    def on_actionSave_As_triggered(self):
        self.save_project(None)
    
    @pyqtSlot()
    def on_actionSave_triggered(self):
        self.save_project(self.PROJECT_FILE)
    
    @pyqtSlot()
    def on_actionImport_to_Project_triggered(self):
        fname = self.getfile(TYPES_IMPORT, ROLE_GLOBAL)
        log("importing file %s" % fname)
        if fname:
            try:
                pj = pystratIO.open_file(fname)
                self.add_recent_file(fname)
                self.merge_project_to_current(pj)
            except Exception as e:
                traceback.print_exc()
                self.error("Error importing file: %s" % str(e))
    
    @pyqtSlot()
    def on_actionPaste_triggered(self):
        cb = QApplication.clipboard()
        text = cb.text()
        if text.strip():
            try:
                pj = pystratIO.open_from_text_data(text)
                self.merge_project_to_current(pj)
            except Exception as e:
                traceback.print_exc()
                self.error("Error importing text from clipboard: %s" % str(e))
        else:
            self.error("No text on clipboard to import.")
    
    def _ask_overwrite_dataset(self, question):
        return QMessageBox.question(self, "Import Dataset", question, 
                                    buttons=QMessageBox.YesToAll|QMessageBox.Yes|QMessageBox.No|QMessageBox.NoToAll|QMessageBox.Cancel, 
                                    defaultButton=QMessageBox.NoButton)
    
    def merge_project_to_current(self, new_project):
        pj = self.PROJECT
        for location_id, location in new_project.locations.items():
            if location_id in pj.locations:
                #attempt to merge
                current_loc = pj.locations[location_id]
                vdata_to_add = []
                skipask = False
                skipask_overwrite = False
                for vdata_id, vdata in location.vdata.items():
                    is_agedepth = (vdata_id == pystratproj.VDATA_ID_AGEDEPTH)
                    
                    if vdata_id in current_loc.vdata:
                        
                        if not skipask:
                            if is_agedepth:
                                question = "Overwrite age-depth model for location %s?" % current_loc.name()
                            else:
                                question = "Overwrite dataset %s in location %s?" % (vdata.name(), current_loc.name())
                                
                            response = self._ask_overwrite_dataset(question)
                            if response == QMessageBox.Cancel:
                                return
                            elif response == QMessageBox.Yes or response == QMessageBox.YesToAll:
                                vdata_to_add.append(vdata)
                                if is_agedepth:
                                    current_loc.age_depth_model = location.age_depth_model
                                    self.set_project_dirty()
                                if response == QMessageBox.YesToAll:
                                    skipask = True
                                    skipask_overwrite = True
                            elif response == QMessageBox.NoToAll:
                                skipask = True
                                skipask_overwrite = False
                        else:
                            if skipask_overwrite:
                                vdata_to_add.append(vdata)
                                if is_agedepth:
                                    current_loc.age_depth_model = location.age_depth_model
                                    self.set_project_dirty()
                    else:
                        vdata_to_add.append(vdata)
                        if is_agedepth:
                            current_loc.age_depth_model = location.age_depth_model
                            self.set_project_dirty()
                
                for vdata in vdata_to_add:
                    current_loc.add_vdata(vdata)
                    self.set_project_dirty()
            else:
                #add location and be gone!
                pj.add_location(location)
                self.set_project_dirty()
        self.refresh_trees()
        
    def set_project_dirty(self):
        log("set project dirty")
        self.PROJECT_DIRTY = True
        self.setWindowModified(True)
    
    def load_file(self, fname):
        log("loading file %s" % fname)
        if fname:
            try:
                pj = pystratIO.open_file(fname)
                self.set_project(pj)
                self.add_recent_file(fname)
                #set dirty if imported by CSV
                if fname.endswith(".csv"):
                    self.set_project_dirty()
                    self.PROJECT_FILE = None
                else:
                    self.PROJECT_FILE = fname
                PREFS[PREF_LAST_OPEN_FILE] = fname
            except Exception as e:
                traceback.print_exc()
                self.error("Error opening file: %s" % str(e))
    
    def save_project(self, fname=None):
        print("Saving project with filename %s" % fname)
        if not fname:
            fname = self.putfile(TYPES_SAVE, ROLE_GLOBAL)
        if fname:
            try:
                pystratIO.save_file(self.PROJECT, fname)
                self.PROJECT_FILE = fname
                self.add_recent_file(fname)
                return True
            except Exception as e:
                traceback.print_exc()
                self.error("Error saving file: %s" % str(e))
                return False
        else:
            #do nothing, file dialog cancelled
            return None
    
    def ask_save_changes(self):
        if self.PROJECT_DIRTY:
            result = QMessageBox.question(self, "Save Changes", "Save changes to current project?", 
                                        buttons=QMessageBox.Save|QMessageBox.Discard|QMessageBox.Cancel, 
                                        defaultButton=QMessageBox.NoButton)
            if result == QMessageBox.Save:
                if self.PROJECT_FILE:
                    if self.save_project(self.PROJECT_FILE):
                        return True
                    else:
                        self.error("Could not save project! Try Save As...")
                        return False
                else:
                    result = self.save_project(self.PROJECT_FILE)
                    if result is True:
                        return True
                    elif result is None:
                        return False #no file, dialog was cancelled
                    else:
                        self.error("Could not save project! Try a different filename.")
                        return False
            elif result == QMessageBox.Discard:
                return True
            else:
                return False
        else:
            return True
    
    def add_recent_file(self, fname):
        recent_files = PREFS[PREF_RECENT_FILES]
        if recent_files is None:
            recent_files = []
        elif fname in recent_files:
            recent_files.remove(fname)
        recent_files.append(fname)
        
        log("setting project to clean")
        #set window title
        self.setWindowTitle("PyStrat [%s]" % os.path.basename(fname))
        #set window to be clean
        self.PROJECT_DIRTY = False
        self.setWindowModified(False)
        
        if len(recent_files) > 10:
            recent_files.pop(0)
        self.refresh_recent_files()
        
    def refresh_recent_files(self):
        files_menu = self.menu_RecentFiles
        files_menu.clear()
        for file_name in reversed(PREFS[PREF_RECENT_FILES]):
            action = self.create_action(file_name, functools.partial(self.load_file_if_unmodified, new_fname=file_name))
            files_menu.addAction(action)
    
    def load_file_if_unmodified(self, new_fname):
        if self.ask_save_changes():
            self.load_file(new_fname)    
    
    def set_project(self, pystrat_project):
        if pystrat_project is None:
            pystrat_project = pystratproj.PyStratProject()
            del PREFS[PREF_LAST_OPEN_FILE]
        self.clear_data_table()
        self.PROJECT = pystrat_project
        self.refresh_trees()
        
        self.PROJECT_DIRTY = False
        self.setWindowModified(False)
        
    
    #===============VDATA AND LOCATION LIST METHODS=================#
    
    ROLE_LOCATION_ID = Qt.UserRole
    ROLE_VDATA_ID = Qt.UserRole+1
    
    def setup_trees(self):
        self.active_tree = self.tree_location
        self.refresh_trees()
    
    def tree_item_data(self, item):
        vdata_id = item.data(0, PyStratMainWindow.ROLE_VDATA_ID)
        location_id = item.data(0, PyStratMainWindow.ROLE_LOCATION_ID)
        return vdata_id, location_id
    
    @pyqtSlot(int)
    def on_tab_trees_currentChanged(self, index):
        if index == 0:
            self.active_tree = self.tree_location
        else:
            self.active_tree = self.tree_vdata
            
    @pyqtSlot(QModelIndex)
    def on_tree_location_doubleClicked(self, model_index):
        if model_index:
            item = self.tree_location.itemFromIndex(model_index)
            if item:
                #enable only if has both location and vdata id
                vdata_id, location_id = self.tree_item_data(item)
                if vdata_id and location_id:
                    vdata_list = self.tree_item_to_vdata(item, True)
                    self.data_table_add_vdata(vdata_list)
    
    @pyqtSlot(QModelIndex)
    def on_tree_vdata_doubleClicked(self, model_index):
        if model_index:
            item = self.tree_vdata.itemFromIndex(model_index)
            if item:
                #enable only if has both location and vdata id
                vdata_id, location_id = self.tree_item_data(item)
                if vdata_id and location_id:
                    vdata_list = self.tree_item_to_vdata(item, True)
                    self.data_table_add_vdata(vdata_list)
    
    @pyqtSlot(QPoint)
    def on_tree_vdata_customContextMenuRequested(self, point):
        items = self.tree_vdata.selectedItems()
        if len(items) > 0:
            vdata_list = []
            mode = None
            for item in items:
                vdata_id, location_id = self.tree_item_data(item)
                if location_id and vdata_id:
                    if not mode is None and not mode == "single_vdata":
                        return None
                    if mode is None:
                        mode = "single_vdata"
                elif vdata_id and not location_id:
                    if not mode is None and not mode == "group_vdata":
                        return None
                    if mode is None:
                        mode = "group_vdata"
                vdata_from_item = self.tree_item_to_vdata(item, False)
                for vdata in vdata_from_item:
                    vdata_list.append(vdata)
            
            if mode == "group_vdata":
                menu = self.generate_vdata_context_menu(vdata_list)
            elif mode == "single_vdata":
                menu = self.generate_single_vdata_context_menu(vdata_list)
            else:
                return None
            menu.exec_(self.tree_vdata.mapToGlobal(point))

    
    @pyqtSlot(QPoint)
    def on_tree_location_customContextMenuRequested(self, point):
        items = self.tree_location.selectedItems()
        if len(items) > 0:
            locations = []
            vdata_list = []
            
            for item in items:
                vdata_id, location_id = self.tree_item_data(item)
                if vdata_id and location_id:
                    if locations:
                        return None
                    item_vdata = self.tree_item_to_vdata(item, False)
                    for vdata in item_vdata:
                        vdata_list.append(vdata)
                elif location_id:
                    if vdata_list:
                        return None
                    locations.append(self.PROJECT.locations[location_id])
            
            if locations:
                menu = self.generate_location_context_menu(locations)
            elif vdata_list:
                menu = self.generate_single_vdata_context_menu(vdata_list)
            else:
                return None
            menu.exec_(self.tree_location.mapToGlobal(point))
    
    @pyqtSlot(QTreeWidgetItem)
    def on_tree_vdata_itemExpanded(self, item):
        self._set_vdata_expanded(item, True)
    
    @pyqtSlot(QTreeWidgetItem)
    def on_tree_vdata_itemCollapsed(self, item):
        self._set_vdata_expanded(item, False)    
    
    def _set_vdata_expanded(self, item, state):
        vdata_id = item.data(0, PyStratMainWindow.ROLE_VDATA_ID)
        for location in self.PROJECT.locations.values():
            if vdata_id in location.vdata:
                location.vdata[vdata_id].set_metadata(KEY_TREE_EXPANDED, state)
    
    @pyqtSlot(QTreeWidgetItem)
    def on_tree_location_itemExpanded(self, item):
        location_id = item.data(0, PyStratMainWindow.ROLE_LOCATION_ID)
        self.PROJECT.locations[location_id].set_metadata(KEY_TREE_EXPANDED, True)

    @pyqtSlot(QTreeWidgetItem)
    def on_tree_location_itemCollapsed(self, item):
        location_id = item.data(0, PyStratMainWindow.ROLE_LOCATION_ID)
        self.PROJECT.locations[location_id].set_metadata(KEY_TREE_EXPANDED, False)
        
    def set_vdata_selected(self, vdata_list):
        tree = self.active_tree
        for vdata in vdata_list:
            for top_index in range(tree.topLevelItemCount()):
                top_item = tree.topLevelItem(top_index)
                for child_index in range(top_item.childCount()):
                    child_item = top_item.child(child_index)
                    vdata_id, location_id = self.tree_item_data(child_item)
                    if loc_id(vdata)==location_id and vdata._id == vdata_id:
                        child_item.setSelected(True)
                        top_item.setExpanded(True)
        
    def generate_vdata_context_menu(self, vdata_list):
        cmenu = QMenu()
        locs = set()
        unit_set = set()
        vdata_ids = set()
        for vdata in vdata_list:
            locs.add(loc_id(vdata))
            if not vdata.value_unit is None:
                unit_set.add(vdata.value_unit)
            vdata_ids.add(vdata._id)
        
        num_vdata = len(vdata_list)
        vdata_label = "dataset" if num_vdata == 1 else "datasets"
        status_action = self.create_action("%s %s selected" % (num_vdata, vdata_label))
        status_action.setEnabled(False)
        cmenu.addAction(status_action)
        cmenu.addAction(self.create_separator())    
        
        cmenu.addAction(self.create_action("Add to data table", self.on_button_sendTreeToData_clicked))
        cmenu.addAction(self.create_action("Autochart..", functools.partial(self.autochart_launch, vdata_list=vdata_list)))
        cmenu.addAction(self.create_action("Standardized Anomalies", functools.partial(self.standardized_anomalies, vdata_list=vdata_list)))
        
        if len(vdata_ids) == 1:
            cmenu.addAction(self.create_action("Rename..", functools.partial(self.rename_vdata, vdata=vdata_list[0])))
        if len(unit_set) == 1:
            convert_menu = QMenu("Convert unit to", cmenu)
            current_unit = unit_set.pop()
            self.generate_units_menu(convert_menu, vdata_list, units.get_category(vdata_list[0].value_unit), current_unit)
            cmenu.addMenu(convert_menu)
        set_menu = QMenu("Force unit to", cmenu)
        self.generate_units_menu(set_menu, vdata_list, None)
        cmenu.addMenu(set_menu)
            
        copy_menu = QMenu("Copy Text", cmenu)
        copy_all_menu = QMenu("All", copy_menu)
        self.generate_copy_vdata_text_menu(copy_all_menu, pystratIO.EXPORT_MODE_ALL, vdata_list)
        copy_menu.addMenu(copy_all_menu)
        copy_averages_menu = QMenu("Averages", copy_menu)
        self.generate_copy_vdata_text_menu(copy_averages_menu, pystratIO.EXPORT_MODE_SUMMARY, vdata_list)
        copy_menu.addMenu(copy_averages_menu)
        cmenu.addMenu(copy_menu)
        
        cmenu.addAction(self.create_action("Delete", functools.partial(self.delete_vdata, vdata_list=vdata_list)))
        
        cmenu.addAction(self.create_separator())
        cmenu.addAction(self.create_action("Properties"))
        
        return cmenu
    
    def generate_location_context_menu(self, location_list):
        cmenu = QMenu(self)
        
        num_locs = len(location_list)
        loc_label = "location" if num_locs == 1 else "locations"
        status_action = self.create_action("%s %s selected" % (num_locs, loc_label))
        status_action.setEnabled(False)
        cmenu.addAction(status_action)
        cmenu.addAction(self.create_separator())
        
        cmenu.addAction(self.create_action("Add to data table", self.on_button_sendTreeToData_clicked))
        vdata_list = []
        for location in location_list:
            for vdata in location.vdata.values():
                vdata_list.append(vdata)
        cmenu.addAction(self.create_action("Autochart..", functools.partial(self.autochart_launch, vdata_list=vdata_list)))
        
        if len(location_list) == 1:
            cmenu.addAction(self.create_action("Rename..", functools.partial(self.rename_location, location=location_list[0])))
            cmenu.addAction(self.create_action("Import Age-Depth Model", functools.partial(self.age_depth_import, location=location_list[0])))
        
        copy_menu = QMenu("Copy Text", cmenu)
        copy_all_menu = QMenu("All", copy_menu)
        self.generate_copy_vdata_text_menu(copy_all_menu, pystratIO.EXPORT_MODE_ALL, vdata_list)
        copy_menu.addMenu(copy_all_menu)
        copy_averages_menu = QMenu("Averages", copy_menu)
        self.generate_copy_vdata_text_menu(copy_averages_menu, pystratIO.EXPORT_MODE_SUMMARY, vdata_list)
        copy_menu.addMenu(copy_averages_menu)
        cmenu.addMenu(copy_menu)
        
        cmenu.addAction(self.create_action("Delete", functools.partial(self.delete_locations, location_list=location_list)))
        cmenu.addAction(self.create_separator())
        cmenu.addAction(self.create_action("Properties"))
        
        return cmenu
    
    def generate_single_vdata_context_menu(self, vdata_list, menu=None):
        if not menu:
            cmenu = QMenu(self)
        else:
            cmenu = menu
        
        dendrograms = []
        for vdata in vdata_list:
            if isinstance(vdata, pystratmodel.Dendrogram):
                dendrograms.append(vdata)
        
        num_vdata = len(vdata_list)
        vdata_label = "dataset" if num_vdata == 1 else "datasets"
        status_action = self.create_action("%s %s selected" % (num_vdata, vdata_label))
        status_action.setEnabled(False)
        cmenu.addAction(status_action)
        cmenu.addAction(self.create_separator())
        
        cmenu.addAction(self.create_action("Add To Data Table", self.on_button_sendTreeToData_clicked))
        cmenu.addAction(self.create_action("Autochart..", functools.partial(self.autochart_launch, vdata_list=vdata_list)))
        cmenu.addAction(self.create_action("Standardized Anomalies", functools.partial(self.standardized_anomalies, vdata_list=vdata_list)))

        if len(dendrograms) == 0:
            cmenu.addAction(self.create_action("CONISS Constrained Cluster", functools.partial(self.cluster_vdata, vdata_list=vdata_list)))
        
        copy_menu = QMenu("Copy Text", cmenu)
        copy_all_menu = QMenu("All", copy_menu)
        self.generate_copy_vdata_text_menu(copy_all_menu, pystratIO.EXPORT_MODE_ALL, vdata_list)
        copy_menu.addMenu(copy_all_menu)
        copy_averages_menu = QMenu("Averages", copy_menu)
        self.generate_copy_vdata_text_menu(copy_averages_menu, pystratIO.EXPORT_MODE_SUMMARY, vdata_list)
        copy_menu.addMenu(copy_averages_menu)
        cmenu.addMenu(copy_menu)
        
        copyto_menu = QMenu("Copy To", cmenu)
        self.generate_reassign_menu(copyto_menu, True, vdata_list)
        cmenu.addMenu(copyto_menu)
        
        moveto_menu = QMenu("Move To", cmenu)
        self.generate_reassign_menu(moveto_menu, False, vdata_list)
        cmenu.addMenu(moveto_menu)
        
        cmenu.addAction(self.create_action("Delete", functools.partial(self.delete_vdata, vdata_list=vdata_list)))
        cmenu.addAction(self.create_separator())
        cmenu.addAction(self.create_action("Properties"))
        return cmenu
    
    def generate_units_menu(self, umenu, vdata_list, category, current_unit=None):
        if category:
            units_text = units.category[category]
            for unit in units_text:
                unit_action = self.create_action(units.label(unit), functools.partial(self.units_convert, vdata_list=vdata_list, unit=unit))
                if unit == current_unit:
                    unit_action.setEnabled(False)
                umenu.addAction(unit_action)
        else:
            umenu.addAction(self.create_action("Reset unit to none", functools.partial(self.units_convert, vdata_list=vdata_list, unit=None)))
            for category_text in units.categories:
                submenu = QMenu(category_text, umenu)
                units_text = units.category[category_text]
                for unit in units_text:
                    submenu.addAction(units.label(unit), functools.partial(self.units_force, vdata_list=vdata_list, unit=unit))
                umenu.addMenu(submenu)
    
    def generate_copy_vdata_text_menu(self, cmenu, mode, vdata_list):
        copy_spreadsheet = self.create_action("Copy Spreadsheet", functools.partial(self.copy_vdata_text, vdata_list=vdata_list, mode=mode, export_format=pystratIO.FORMAT_TSV))
        copy_r = self.create_action("Copy R Matrix", functools.partial(self.copy_vdata_text, vdata_list=vdata_list, mode=mode, export_format=pystratIO.FORMAT_R_MATRIX))
        cmenu.addAction(copy_spreadsheet)
        cmenu.addAction(copy_r)
    
    def generate_reassign_menu(self, cmenu, copy, vdata_list):
        locations = sorted(self.PROJECT.locations.values(), key=lambda loc: loc.name())
        locations.append(None)
        for location in locations:
            name = location.name() if location else "New Location.."
            act = self.create_action(name, functools.partial(self.reassign_vdata, copy=copy, vdata_list=vdata_list, location=location))
            cmenu.addAction(act)
        
            
    def units_force(self, vdata_list, unit):
        for vdata in vdata_list:
            vdata.value_unit = unit
        self.set_project_dirty()
        self.refresh_trees() 
            
    def units_convert(self, vdata_list, unit):
        new_list = []
        for vdata in vdata_list:
            new_vdata = dataprocessing.convert_units(vdata, unit)
            new_list.append(new_vdata)
            self.PROJECT.locations[loc_id(vdata)].add_vdata(new_vdata)
        self.set_project_dirty()
        self.refresh_trees()
        self.set_vdata_selected(new_list)
    
    def cluster_vdata(self, vdata_list):
        try:
            locs = set()
            for vdata in vdata_list:
                locs.add(loc_id(vdata))
            if len(locs) > 1:
                raise ValueError("Cannot cluster for more than one location (yet!)")
            locs = list(locs)
            dendrogram = dataprocessing.cluster_vdata(vdata_list, True)
            self.PROJECT.locations[locs[0]].add_vdata(dendrogram)
            self.set_project_dirty()
            self.refresh_trees()
            self.set_vdata_selected([dendrogram,])
            
            vdata_list.append(dendrogram)
            self.autochart_launch(vdata_list)
        except Exception as e:
            traceback.print_exc()
            self.error("Error performing cluster analysis: %s" % str(e))
    
    def standardized_anomalies(self, vdata_list):
        new_vdata = []
        try:
            for vdata in vdata_list:
                new_vdata.append(dataprocessing.standardized_anomalies(vdata))
        except Exception as e:
            traceback.print_exc()
            self.error("Error calculating standardized anomalies: %s" % str(e))
            return
            
        if new_vdata:
            for vdata in new_vdata:
                self.PROJECT.locations[loc_id(vdata)].add_vdata(vdata)
            self.set_project_dirty()
            self.refresh_trees()
            self.set_vdata_selected(new_vdata)
    
    def copy_vdata_text(self, vdata_list, export_format, mode):
        try:
            text = pystratIO.export_vdata(vdata_list, filename=None, file_format=export_format, mode=mode)
            cb = QApplication.clipboard()
            cb.setText(text)
        except Exception as e:
            traceback.print_exc()
            self.error("Error converting data to text: %s" % str(e))
    
    def rename_location(self, location):
        title = "Rename Location: %s" % location.name()
        rawname = location.metadata(pystratmodel.METADATA_NAME)
        message = "Current name: %s\nCurrent id: %s" % (rawname, location._id)
        newname, ok = QInputDialog.getText(self, title, message, text=rawname)
        if ok:
            if newname and newname.strip():
                newname = newname.strip()
                for existing_location in self.PROJECT.locations.values():
                    if existing_location.metadata(pystratmodel.METADATA_NAME) == newname or existing_location.name() == newname:
                        self.error("Name %s already exists." % newname)
                        return
                location.set_name(newname)
                self.set_project_dirty()
                self.refresh_trees()
            else:
                self.error("Must enter non-empty name.")
    
    def rename_vdata(self, vdata): #NEEDS HELP names and IDs are important to distinguish
        title = "Rename Dataset: %s" % vdata.name()
        rawname = vdata.metadata(pystratmodel.METADATA_NAME)
        message = "Current name: %s\nCurrent id: %s" % (rawname, vdata._id)
        newname, ok = QInputDialog.getText(self, title, message, text=rawname)
        if ok:
            if newname and newname.strip():
                newname = newname.strip()
                name_plus_unit = "%s (%s)" % (newname, units.label(vdata.value_unit)) #a bit of a hack, could change if naming strategy is different
                new_id = name_plus_unit.lower().replace(" ", "_") #bit of a hack, method for changing names to Ids from CSV could change
                location = self.PROJECT.locations[loc_id(vdata)]
                for existing_vdata in location.vdata.values():
                    if existing_vdata.metadata(pystratmodel.METADATA_NAME) == newname or existing_vdata.name() == newname or existing_vdata._id == new_id:
                        self.error("%s already exists in location %s." % (newname, location.name()))
                        return
                vdata.set_name(newname)
                self.set_project_dirty()
                self.refresh_trees()
            else:
                self.error("Must enter non-empty name.")
    
    def reassign_vdata(self, vdata_list, copy, location):
        if not location:
            title = "New Location"
            message = "Enter a name for new location:"
            newname, ok = QInputDialog.getText(self, title, message, text="New Location")
            if ok:
                if newname and newname.strip():
                    newname = newname.strip()
                    for existing_location in self.PROJECT.locations.values():
                        if existing_location.metadata(pystratmodel.METADATA_NAME) == newname or existing_location.name() == newname:
                            self.error("Name %s already exists." % newname)
                            return
                    location = pystratproj.DataLocation(_id=newname)
                    location.set_name(newname)
                    self.PROJECT.add_location(location)
                else:
                    self.error("Must enter non-empty name.")
                    return
            else:
                return
        
        new_vdatas = []    
        for vdata in vdata_list:
            old_location = self.PROJECT.locations[loc_id(vdata)]
            if not copy:
                new_vdata = vdata
            else:
                new_vdata = vdata.copy()
            
            if new_vdata._id in location.vdata:
                if not self.question("Overwrite %s in location %s?" % (new_vdata.name(), location.name())):
                    continue
            location.add_vdata(new_vdata)
            new_vdatas.append(new_vdata)
            if not copy:
                del old_location.vdata[vdata._id]
                
            
        self.set_project_dirty()
        self.refresh_trees()
        self.set_vdata_selected(new_vdatas)
                    
    
    def delete_vdata(self, vdata_list):
        message = "Delete dataset %s? This cannot be undone!" % vdata_list[0].name() if len(vdata_list) == 1 else "Delete %s datasets? This cannot be undone!" % len(vdata_list)
        if self.question(message):
            for vdata in vdata_list:
                del self.PROJECT.locations[loc_id(vdata)].vdata[vdata._id]
            self.refresh_trees()
            self.refresh_data_table()
            self.set_project_dirty()
    
    def delete_locations(self, location_list):
        message = "Delete location %s including all associated data? This cannot be undone!" % location_list[0].name() if len(location_list) == 1 else "Delete %s locations? This cannot be undone!" % len(location_list)
        if self.question(message):
            for location in location_list:
                del self.PROJECT.locations[location._id]
            self.refresh_trees()
            self.refresh_data_table()
            self.set_project_dirty()
    
    
    def age_depth_import(self, location):
        csvfile = self.getfile("Comma Separated Values (*.csv)", ROLE_GLOBAL)
        if csvfile:
            try:
                proj = pystratIO.open_file(csvfile)
                models = pystratIO.extract_age_depth_models(proj)
                if len(models) > 1:
                    raise ValueError("More than 1 age depth model found in file")
                elif len(models) == 0:
                    raise ValueError("No age depth models found in file")
                adm = list(models.values())[0]
                location.age_depth_model = adm
                ad_vdata = adm.to_vdata()
                ad_vdata._id = pystratproj.VDATA_ID_AGEDEPTH
                ad_vdata.set_name("Age Depth Model")
                location.add_vdata(ad_vdata)
                
                self.refresh_trees()
                self.set_project_dirty()
            except Exception as e:
                traceback.print_exc()
                self.error("Error importing age-depth model: %s" % str(e))
                
    
    def refresh_trees(self):
        if self.PROJECT:
            pj = self.PROJECT
            
            location_names = []
            location_ids = []
            expanded_loc = []
            for location_id, location in pj.locations.items():
                name = location.name()
                location_names.append(name)
                expanded_loc.append(location.metadata(KEY_TREE_EXPANDED))
                location_ids.append(location_id)
            
            vdata_names = []
            vdata_ids = []
            
            vdata_names_dict = {}
            
            for location_id in location_ids:
                names = []
                ids = []
                loc = pj.locations[location_id]
                for vdata_id, vdata in loc.vdata.items():
                    name = vdata.name()
                    names.append(name)
                    ids.append(vdata_id)
                    vdata_names_dict[vdata_id] = name
                    
                vdata_names.append(names)
                vdata_ids.append(ids)
                    
            self.set_tree_values(self.tree_location, location_names, expanded_loc, vdata_names, location_ids, vdata_ids, 
                                 PyStratMainWindow.ROLE_LOCATION_ID, PyStratMainWindow.ROLE_VDATA_ID, Qt.AscendingOrder)
            
            vdata_names = []
            vdata_ids = []
            expanded_vdata = []
            location_names = []
            location_ids = []
            for vdata_id, vdata_name in vdata_names_dict.items():
                vdata_names.append(vdata_name)
                vdata_ids.append(vdata_id)
                is_expanded = False
                names = []
                ids = []
                for location_id, location in pj.locations.items():
                    if vdata_id in location.vdata:
                        ids.append(location_id)
                        names.append(location.name())
                        is_expanded = location.vdata[vdata_id].metadata(KEY_TREE_EXPANDED)
                location_names.append(names)
                location_ids.append(ids)
                expanded_vdata.append(is_expanded)
            
            self.set_tree_values(self.tree_vdata, vdata_names, expanded_vdata, location_names, vdata_ids, location_ids, 
                                 PyStratMainWindow.ROLE_VDATA_ID, PyStratMainWindow.ROLE_LOCATION_ID, Qt.AscendingOrder)    
        else:
            self.tree_location.clear()
            self.tree_vdata.clear()
        
        
    
    
    def set_tree_values(self, tree_widget, top_level_names, top_level_expanded, sub_item_names, 
                        top_level_ids, sub_item_ids, top_level_role, sub_item_role, 
                        name_order=Qt.AscendingOrder):
        
        tree_widget.clear()
        for top_level_index in range(len(top_level_names)):
            names = sub_item_names[top_level_index]
            ids = sub_item_ids[top_level_index]
            top_level_item = QTreeWidgetItem(tree_widget)
            top_level_item.setText(0, top_level_names[top_level_index])
            top_level_item.setData(0, top_level_role, top_level_ids[top_level_index])
            if top_level_expanded[top_level_index]:
                top_level_item.setExpanded(True)
            else:
                top_level_item.setExpanded(False)
            
            for sub_item_index in range(len(names)):
                sub_item = QTreeWidgetItem(top_level_item)
                sub_item.setText(0, names[sub_item_index])
                sub_item.setData(0, top_level_role, top_level_ids[top_level_index])
                sub_item.setData(0, sub_item_role, ids[sub_item_index])
                top_level_item.addChild(sub_item)
                
            top_level_item.sortChildren(0, name_order)
            tree_widget.addTopLevelItem(top_level_item)
        tree_widget.sortItems(0, name_order)
    
    @pyqtSlot()
    def on_button_sendTreeToData_clicked(self):
        selected_items = self.active_tree.selectedItems()
        vdata_list = []
        for item in selected_items:
            for vdata in self.tree_item_to_vdata(item, True):
                if isinstance(vdata, pystratmodel.VerticalData):
                    vdata_list.append(vdata)
        self.data_table_add_vdata(vdata_list)
        
    def tree_item_to_vdata(self, item, vdata_only):
        pj = self.PROJECT
        vdata_list = []
        vdata_id = item.data(0, PyStratMainWindow.ROLE_VDATA_ID)
        location_id = item.data(0, PyStratMainWindow.ROLE_LOCATION_ID)
        if vdata_id and location_id:
            try:
                log("adding vdata %s for location %s" %(vdata_id, location_id))
                vdata = pj.locations[location_id].vdata[vdata_id]
                if isinstance(vdata, pystratmodel.VerticalData) or not vdata_only:
                    vdata_list.append(vdata)
            except KeyError:
                self.error("Could not find key % or %s" % (vdata_id, location_id))
        elif vdata_id:
            try:
                log("adding vdata %s for all locations" % vdata_id)
                for location_id, location in pj.locations.items():
                    if vdata_id in location.vdata:
                        vdata = location.vdata[vdata_id]
                        if isinstance(vdata, pystratmodel.VerticalData) or not vdata_only:
                            vdata_list.append(vdata)
            except KeyError:
                self.error("Could not find key when adding vdata %s for multiple locations" % vdata_id)
        elif location_id:
            try:
                log("adding all vdata for location %s" % location_id)
                for vdata in sorted(pj.locations[location_id].vdata.values(), key=lambda it: it.name()):
                    if isinstance(vdata, pystratmodel.VerticalData) or not vdata_only:
                        vdata_list.append(vdata)
            except KeyError:
                self.error("Could not find key when adding all vdata for location %s" % location_id)
        
        return vdata_list
                
     
    #==================AUTOCHART====================================#
    
    def autochart_launch(self, vdata_list):
        if not vdata_list:
            vdata_list = []
        self.autochartDialog.set_vdata(self.PROJECT, vdata_list)
             
    
    #==================DATA TABLE METHODS===========================#
    
    def setup_data_table(self):
        self.data_table_vdata = []
        self.data_table_editable = False
        self.table_data.addAction(self.actionCopy)
        self.refresh_data_table()
        
        self.combo_dataTableMode.addItem("Values")
        self.combo_dataTableMode.addItem("Number of readings (n)")
        self.combo_dataTableMode.addItem("Error")
        self.combo_dataTableMode.addItem("Sample Ids")
        
    @pyqtSlot()
    def on_button_dataTableClear_clicked(self):
        if self.question("Clear data table?"):
            self.clear_data_table()
    
    def on_button_clearSelection_clicked(self):
        self.table_data.clearSelection()
    
    @pyqtSlot()
    def on_button_autoChart_clicked(self):
        #get selected columns TODO possibly only do this if selected columns are the whole column
        if len(self.data_table_vdata) > 0:
            selection_ranges = self.table_data.selectedRanges()
            if len(selection_ranges) > 0:
                vdata_out = []
                for sel_range in selection_ranges:
                    for col in range(sel_range.leftColumn(), sel_range.rightColumn()+1):
                        if col >= 1:
                            vdata_out.append(self.data_table_vdata[col-1])
                self.autochart_launch(vdata_out)
            else:
                #export all
                self.autochart_launch(self.data_table_vdata)
        else:
            self.error("Add data to the data table to use AutoChart.")
    
    @pyqtSlot()
    def on_menuData_Table_aboutToShow(self):
        self.menuData_Table.clear()
        self.generate_single_vdata_context_menu(self.data_table_vdata, menu=self.menuData_Table)
    
    @pyqtSlot()
    def on_actionCopy_triggered(self):
        self.data_table_copy_selected()
    
    def data_table_add_vdata(self, vdata_list):
        for new_vdata in vdata_list:
            if new_vdata in self.data_table_vdata:
                log("not adding duplicate vdata %s" % new_vdata._id)
            else:
                log("adding vdata with id %s" % new_vdata._id)
                self.data_table_vdata.append(new_vdata)
        self.refresh_data_table()
    
    def data_table_copy_selected(self):
        selection_ranges = self.table_data.selectedRanges()
        if len(selection_ranges) > 0:
            #only reading one selection range
            selection_range = selection_ranges[0]
            nl = ""
            string_out = ""
            dt = self.table_data
            for row in range(selection_range.topRow(), selection_range.bottomRow()+1):
                string_row = nl + ""
                tab = ""
                for col in range(selection_range.leftColumn(), selection_range.rightColumn()+1):
                    val = dt.item(row, col).text()
                    string_row += tab + val
                    tab = "\t"
                nl = "\n"
                
                string_out += string_row
            if string_out:
                if selection_range.rowCount() == dt.rowCount():
                    header_string = ""
                    tab = ""
                    for col in range(selection_range.leftColumn(), selection_range.rightColumn()+1):
                        val = dt.horizontalHeaderItem(col).text()
                        header_string += tab + val
                        tab = "\t"
                    string_out = header_string + "\n" + string_out
                
                cb = QApplication.clipboard()
                cb.setText(string_out)
            
        else:
            pass
    
    def clear_data_table(self):
        self.data_table_vdata = []
        self.refresh_data_table()
    
    def refresh_data_table(self):
        log("refreshing data table")
        if self.data_table_vdata:
            pj = self.PROJECT
            
            #ensure all vdata/locations still exist
            toremove = []
            for vdata in self.data_table_vdata:
                try:
                    exists = vdata._id in pj.locations[loc_id(vdata)].vdata
                    if not exists:
                        toremove.append(vdata)
                except KeyError:
                    toremove.append(vdata)
            for vdata in toremove:
                self.data_table_vdata.remove(vdata)
            
            depths = set()
            for vdata in self.data_table_vdata:
                depths = depths.union(set(vdata.depths()))
            depths_list = sorted(depths)
            values_all = []
            n_all = []
            error_all = []
            sample_ids_all = []
            tooltips_all = []
            headers = []
            values_all.append(depths_list)
            n_all.append(depths_list)
            error_all.append(depths_list)
            sample_ids_all.append(depths_list)
            tooltips_all.append([None for index in range(len(depths_list))])
            headers.append("Depth") #TODO this will need to be either depth or age
            
            multiple_locations = len(pj.locations) > 1
            
            for vdata in self.data_table_vdata:
                name = vdata.name()
                loc_name = pj.locations[loc_id(vdata)].name()
                values = []
                n_vals = []
                errors = []
                sample_id_vals = []
                
                tooltips = []
                for depth in depths_list:
                    value = vdata.data_summary(depth)
                    if value:
                        value_formatted = format_value(value[0])
                        values.append(value_formatted)
                        n_vals.append(value[1])
                        all_vals, sample_ids = vdata.data_all(depth)
                        if len(sample_ids) > 1:
                            er_amnt = format_error(calculate_error(value[0], value[1], value[2]))
                            errors.append(er_amnt)
                            tooltip = "%s \u00B1 %s" % (value_formatted, er_amnt)
                            for index in range(len(sample_ids)):
                                tooltip += "\n%s: %s" % (sample_ids[index], all_vals[index])
                            tooltips.append(tooltip)
                            sample_id_vals.append(", ".join(sample_ids))
                        else:
                            tooltips.append("%s: %s" % (sample_ids[0], value_formatted))
                            sample_id_vals.append(sample_ids[0])
                            errors.append("0.0")
                    else:
                        values.append(None)
                        tooltips.append(None)
                        n_vals.append(None)
                        errors.append(None)
                        sample_id_vals.append(None)
                if multiple_locations:        
                    headers.append("%s: %s" % (loc_name, name))
                else:
                    headers.append(name)
                values_all.append(values)
                tooltips_all.append(tooltips)
                n_all.append(n_vals)
                error_all.append(errors)
                sample_ids_all.append(sample_id_vals)
            
            self.data_table_values = values_all
            self.data_table_n_vals = n_all
            self.data_table_sample_ids = sample_ids_all
            self.data_table_tooltips = tooltips_all
            self.data_table_errors = error_all
            self.data_table_headers = headers
            
            self.data_table_implicit_refresh()
        else:
            empty = [[]]
            self.data_table_values = empty
            self.data_table_n_vals = empty
            self.data_table_sample_ids = empty
            self.data_table_tooltips = empty
            self.data_table_errors = empty
            self.data_table_headers = []
            
            self.data_table_implicit_refresh()
            
    
    def on_combo_dataTableMode_currentIndexChanged(self, in_object):
        #disable on startup
        if not self.LOADING:
            self.data_table_implicit_refresh()
    
    def data_table_implicit_refresh(self):
        combo_mode_index = self.combo_dataTableMode.currentIndex()
        if combo_mode_index == 0:
            table_vals = self.data_table_values
        elif combo_mode_index == 1:
            table_vals = self.data_table_n_vals
        elif combo_mode_index == 2:
            table_vals = self.data_table_errors
        elif combo_mode_index == 3:
            table_vals = self.data_table_sample_ids
        else:
            table_vals = [[]]
        
        self.data_table_data(self.data_table_headers, table_vals, self.data_table_tooltips)
    
    def data_table_data(self, horizontal_headers, two_dim_array, tool_tips):
        cols = len(two_dim_array)
        if cols > 0:
            rows = len(two_dim_array[0])
        else:
            rows = 0
        self.table_data.setColumnCount(cols)
        self.table_data.setRowCount(rows)
        if horizontal_headers:
            self.table_data.horizontalHeader().show()
            self.table_data.setHorizontalHeaderLabels(horizontal_headers)
        else:
            self.table_data.horizontalHeader().hide()

        replicate_bg_color = QColor(PREFS[PREF_REPLICATE_BG_COLOR])
        

        for col in range(cols):
            for row in range(rows):
                value = two_dim_array[col][row]
                if value is None:
                    string_value = ""
                else:
                    string_value = str(value)
                item = QTableWidgetItem(string_value)
                item.setToolTip(tool_tips[col][row])
                
                #background color for replicates
                if self.data_table_n_vals and col > 0:
                    n = self.data_table_n_vals[col][row]
                    if n and n > 1:
                        item.setBackground(QBrush(replicate_bg_color))
                
                if not self.data_table_editable:
                    item.setFlags(item.flags() & ~Qt.ItemIsEditable)
                self.table_data.setItem(row, col, item)
    
    #=======================UTILITY FUNCTIONS=========================#
    
    def create_action(self, text, trigger_slot=None, shortcut=None, iconRes=None, tip=None, checkable=False):
        action = QAction(text, self)
        if iconRes is not None:
            action.setIcon(QIcon(iconRes))
        if shortcut is not None:
            action.setShortcut(shortcut)
        if tip is not None:
            action.setToolTip(tip)
            action.setStatusTip(tip)
        if trigger_slot is not None:
            action.triggered.connect(trigger_slot)
        if checkable:
            action.setCheckable(True)
        return action
    
    def create_separator(self):
        action = QAction("", self)
        action.setSeparator(True)
        return action
    
    def error(self, message):
        QMessageBox.critical(self, "Error", message)
        
    def message(self, message):
        QMessageBox.information(self, "Message", message)
    
    def question(self, question):
        return QMessageBox.question(self, "Confirm", question, 
                                    buttons=QMessageBox.Yes|QMessageBox.No, 
                                    defaultButton=QMessageBox.NoButton) == QMessageBox.Yes
    
    def _last_file(self):
        recent_files = PREFS[PREF_RECENT_FILES]
        if recent_files:
            return recent_files[-1]
        else:
            return None
        
        
    def getfile(self, types, openrole):
        if openrole in PREFS:
            folder = PREFS[openrole]
        elif ROLE_GLOBAL in PREFS:
            folder = PREFS[ROLE_GLOBAL]
        else:
            folder = ""
        
        lastfile = self._last_file()
        if lastfile:
            folder = os.path.dirname(lastfile)
        
        fname = QFileDialog.getOpenFileName(self, "Choose File", folder, types)
        if fname:
            folder = os.path.dirname(fname[0])
            PREFS[openrole] = folder
        return fname[0]
    
    def putfile(self, types, saverole):
        if saverole in PREFS:
            folder = PREFS[saverole]
        elif ROLE_GLOBAL in PREFS:
            folder = PREFS[ROLE_GLOBAL]
        else:
            folder = ""
        
        lastfile = self._last_file()
        if lastfile:
            folder = os.path.dirname(lastfile)
        
        fname = QFileDialog.getSaveFileName(self, "Save File", folder, types)
        if fname:
            folder = os.path.dirname(fname[0])
            PREFS[saverole] = folder
        return fname[0]

    
    #========PREFERENCES, WINDOW SETTINGS, AND THE CLOSE EVENT HANDLER==========#
    
    def closeEvent(self, event):
        if self.ask_save_changes():
            self.save_preferences()
            self.save_window_settings()
            for tfile in self.TEMPFILES:
                try:
                    os.remove(tfile)
                except OSError:
                    print("could not delete file %s" % tfile)
        else:
            event.ignore()
    
    def restore_preferences(self):
        prefs = QSettings("PyStratUserPrefs")
        tmpDict = prefs.value("PyStratUserPrefs/PrefsDict", QVariant({}))
        for key, value in tmpDict.items():
            PREFS[str(key)] = value
        print(PREFS)
    
    def save_preferences(self):
        prefs = QSettings("PyStratUserPrefs")
        prefs.setValue("PyStratUserPrefs/PrefsDict", QVariant(PREFS))
        del prefs
        
    def restore_window_settings(self):
        settings = QSettings("PyStratMainWindow")
        size = settings.value("PyStratMainWindow/Size", QSize(600, 500))
        position = settings.value("PyStratMainWindow/Position", QPoint(0, 0))
        self.move(position)
        self.resize(size)
        state = settings.value("PyStratMainWindow/State")
        if not state is None:
            self.restoreState(settings.value("PyStratMainWindow/State"))
    
    def save_window_settings(self):
        settings = QSettings("PyStratMainWindow")
        settings.setValue("PyStratMainWindow/Size", QVariant(self.size()))
        settings.setValue("PyStratMainWindow/Position", QVariant(self.pos()))
        settings.setValue("PyStratMainWindow/State", QVariant(self.saveState()))
        del settings
        